#!/usr/bin/env python
# encoding=utf-8
#codeby     道长且阻
#email      ydhcui@suliu.net/QQ664284092
#website    http://www.suliu.net


import time
import json
import re
import sys

#import requests
import settings
import models
from core import scan

class TaskManage(object):
    _taskdict = {}
    _schedule = {}

    def __init__(self,**kw):
        self.kwargs = kw

    @staticmethod
    def gettasks():
        return [{k:v.__doc__.strip()} for k,v in TaskManage._taskdict.items()]


    @classmethod
    def task(self,f):
        '''注册普通任务
        '''
        self._taskdict[f.__name__] = f

    @classmethod
    def schedule(self,t):
        '''注册定时任务
        '''
        def F2(f):
            pass
            #self.schedule[f.__name__] = f
        return F2

    def dotask(self,Q):
        Q.task_state = 'finish'
        Q.finishdate = int(time.time())
        Q.save()

    def addtask(self,Q):
        print(Q.task_name,Q.task_host)
        #self.dotask(Q)
        if Q.task_name in self._taskdict.keys():
            self._taskdict[Q.task_name](Q.task_host)
            Q.task_state = 'finish'
            Q.finishdate = int(time.time())
            Q.save()
        #Q.task_pid = result.task_id,
        #Q.task_desc = f.__doc__

    def loop(self):
        M = models.ScanTask
        fn = (M.task_state=='runing')
        while True:
            query = M.select().where(fn).order_by(M.createdate)
            for Q in query:
                self.addtask(Q)
            #time.sleep(1)

@TaskManage.schedule
def updateportinfo(host,**kwargs):
    '''更新主机端口信息
    '''
    PR = models.PortResult()
    fn = ((PR.name=='http',PR.state=='open')|(PR.port==80))
    query = PR.select().where(fn).order_by(PR.createdate).desc()
    scan.UpdatePortInfo(query).save(models.HttpResult)

@TaskManage.task
def portscan(host,**kwargs):
    '''端口扫描
    '''
    M = models.ScanTask
    H = models.HostResult
    P = models.PortResult
    MP = p.get 
    scan.PortScan(taskid,host,**kwargs).save()


#@TaskManage.task
def pathscan(taskid,host,**kwargs):
    '''路径爆破
    '''
    scan.PathScan(taskid,host,**kwargs).save()


@TaskManage.task
def domainscan(taskid,host,**kwargs):
    '''域名扫描
    '''
    scan.DomainScan(taskid,host,**kwargs).save()

@TaskManage.task
def crawler(host,**kwargs):
    '''路径爬取
    '''
    Crawler(taskid,host).start()

if __name__ == '__main__':
    #TM = TaskManage()
    #TM.loop()


    scan.PortScan('192.168.1.1').save()